import { BinarySearchTree, Node } from "./Tree";
import { Compare, defaultCompare } from "./utils";

enum BalanceFactor {
    UNBALANCED_RIGHT = 1,
    SLIGHTLY_UNBALANCED_RIGHT = 2,
    BALANCED = 3,
    SLIGHTLY_UNBALANCED_LEFT = 4,
    UNBALANCED_LEFT = 5
}

export class AVLTree<T> extends BinarySearchTree<T> {

    getNodeHeight(node: Node<T>) {
        if (node == null) {
            return -1;
        }
        return Math.max(this.getNodeHeight(node.left), this.getNodeHeight(node.right)) + 1;
    }

    getBalanceFactor(node: Node<T>) {
        const heightDifference = this.getNodeHeight(node.left) - this.getNodeHeight(node.right);
        switch (heightDifference) {
            case -2:
                return BalanceFactor.UNBALANCED_RIGHT;
            case -1:
                return BalanceFactor.SLIGHTLY_UNBALANCED_RIGHT;
            case 1:
                return BalanceFactor.SLIGHTLY_UNBALANCED_LEFT;
            case 2:
                return BalanceFactor.UNBALANCED_LEFT;
            default:
                return BalanceFactor.BALANCED;
        }
    }

    public insert(key: T) {
        this.root = this.insertNode(this.root, key);
    }

    protected insertNode(node: Node<T>, key: T) {
        if (node == null) {
            return new Node(key);
        } else if (defaultCompare(key, node.key) === Compare.LESS_THAN) {
            node.left = this.insertNode(node.left, key);
        } else if (defaultCompare(key, node.key) === Compare.BIGGER_THAN) {
            node.right = this.insertNode(node.right, key);
        } else {
            return node;
        }

        const balanceFactor = this.getBalanceFactor(node);
        if (balanceFactor == BalanceFactor.UNBALANCED_LEFT) {
            if (defaultCompare(key, node.left.key) == Compare.LESS_THAN) {
                node = this.rotationLL(node);
            } else {
                return this.rotationLR(node);
            }
        }

        if (balanceFactor == BalanceFactor.UNBALANCED_RIGHT) {
            if (defaultCompare(key, node.right.key) == Compare.BIGGER_THAN) {
                node = this.rotationRR(node);
            } else {
                return this.rotationRL(node);
            }
        }
        return node;
    }

    protected removeNode(node: Node<T>, key: T) {
        if (node == null) {
            return null;
        }
        if (defaultCompare(key, node.key) === Compare.LESS_THAN) {
            node.left = this.removeNode(node.left, key);
        } else if (defaultCompare(key, node.key) === Compare.BIGGER_THAN) {
            node.right = this.removeNode(node.right, key);
        } else {
            if (node.left == null && node.right == null) {
                node = null;
            } else if (node.left == null && node.right != null) {
                node = node.right;
            } else if (node.left != null && node.right == null) {
                node = node.left;
            } else {
                const inOrderSuccessor = this.minNode(node.right);
                node.key = inOrderSuccessor.key;
                node.right = this.removeNode(node.right, inOrderSuccessor.key);
            }
        }

        if (node == null) {
            return node;
        }
        const balanceState = this.getBalanceFactor(node);
        if (balanceState === BalanceFactor.UNBALANCED_LEFT) {
            if (
                this.getBalanceFactor(node.left) === BalanceFactor.BALANCED ||
                this.getBalanceFactor(node.left) === BalanceFactor.SLIGHTLY_UNBALANCED_LEFT
            ) {
                return this.rotationLL(node);
            }
            if (this.getBalanceFactor(node.left) === BalanceFactor.SLIGHTLY_UNBALANCED_RIGHT) {
                return this.rotationLR(node.left);
            }
        }

        if (balanceState === BalanceFactor.UNBALANCED_RIGHT) {
            if (
                this.getBalanceFactor(node.right) === BalanceFactor.BALANCED ||
                this.getBalanceFactor(node.right) === BalanceFactor.SLIGHTLY_UNBALANCED_RIGHT
            ) {
                return this.rotationRR(node);
            }
            if (this.getBalanceFactor(node.right) === BalanceFactor.SLIGHTLY_UNBALANCED_LEFT) {
                return this.rotationRL(node.right);
            }
        }
        return node;
    }

    /**
     * Left left case: rotate right   左-左（LL）：向右的单旋转
     * 这种情况出现于节点的左侧子节点的高度大于右侧子节点的高度时，并且左侧子节点也是平衡或左侧较重的
     *
     *       b                           a
     *      / \                         / \
     *     a   e -> rotationLL(b) ->   c   b
     *    / \                             / \
     *   c   d                           d   e
     *
     * @param node Node<T>
     */
    private rotationLL(node: Node<T>) {
        const temp = node.left;
        node.left = temp.right;
        temp.right = node;
        return temp;
    }

    /**
     * Right right case: rotate left
     * 右-右的情况和左-左的情况相反。它出现于右侧子节点的高度大于左侧子节点的高度，并且右侧子节点也是平衡或右侧较重的
     *
     *     a                              b
     *    / \                            / \
     *   c   b   -> rotationRR(a) ->    a   e
     *      / \                        / \
     *     d   e                      c   d
     *
     * @param node Node<T>
     */
    private rotationRR(node: Node<T>) {
        const temp = node.right;
        node.right = temp.left;
        temp.left = node;
        return temp;
    }

    /**
       * Left right case: rotate left then right
       * @param node Node<T>
       */
    private rotationLR(node: Node<T>) {
        node.left = this.rotationRR(node.left);
        return this.rotationLL(node);
    }

    /**
     * Right left case: rotate right then left
     * @param node Node<T>
     */
    private rotationRL(node: Node<T>) {
        node.right = this.rotationLL(node.right);
        return this.rotationRR(node);
    }

}